#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName :start_browser_mcp_script.py
# @Created  :2025/04/05
# @Author   :批量启动Chromium+MCP实例（先启浏览器，后启MCP）

import os, sys, subprocess, time, socket, threading, argparse, shutil
from pathlib import Path
from playwright.sync_api import sync_playwright

# 导入应用配置
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
from app.config import MCP_SERVERS_COUNT, PROJECT_PATH, DATA_PATH

# TCP端口转发相关常量和函数
BUFSIZE = 32 * 1024  # 单次最大缓冲 32 KiB


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='批量启动Chromium + MCP实例（一一对应）')

    # 主要参数
    parser.add_argument('--count', type=int, default=MCP_SERVERS_COUNT,
                        help=f'要启动的实例数量 (默认: {MCP_SERVERS_COUNT}，来自配置文件)')
    parser.add_argument('--mcp-start-port', type=int, default=8931, help='MCP 起始端口号 (默认: 8931)')
    parser.add_argument('--cdp-start-port', type=int, default=9222, help='Chromium CDP 调试端口起始号 (默认: 9222)')
    parser.add_argument('--start-user', type=int, default=1, help='起始用户标识号 (默认: 1)')
    parser.add_argument('--host', type=str, default='127.0.0.1', help='主机地址 (默认: 127.0.0.1)')
    parser.add_argument('--mcp-version', type=str, default='@playwright/mcp@0.0.44',
                        help='MCP版本 (默认: @playwright/mcp@0.0.44)')
    parser.add_argument('--delay', type=float, default=1.0, help='启动实例之间的延迟(秒) (默认: 1.0)')

    # 窗口布局参数
    parser.add_argument('--window-cols', type=int, default=2, help='窗口列数 (默认: 2)')
    parser.add_argument('--window-rows', type=int, default=2, help='窗口行数 (默认: 2)')
    parser.add_argument('--window-margin', type=int, default=10, help='窗口边距 (默认: 10)')

    # 基础端口参数
    parser.add_argument('--base-browser-port', type=int, default=9222, help='浏览器调试端口基础值 (默认: 9222)')
    parser.add_argument('--base-mcp-port', type=int, default=8931, help='MCP端口基础值 (默认: 8931)')
    parser.add_argument('--base-relay-port', type=int, default=19893, help='转发端口基础值 (默认: 19893)')

    # 控制参数
    parser.add_argument('--is_start_mcp', action='store_true', help='是否启动mcp')
    parser.add_argument('--is_start_cdp', action='store_true', help='是否启动Chromium')
    parser.add_argument('--kill-existing', action='store_true', help='关闭已存在的浏览器进程')

    # 端口转发参数
    parser.add_argument('--forward-port', type=int, default=19893, help='端口转发的监听端口 (默认: 19893)')
    parser.add_argument('--forward-target-port', type=int, default=9222, help='端口转发的目标端口 (默认: 9222)')
    parser.add_argument('--forward-target-host', type=str, default='127.0.0.1',
                        help='端口转发的目标主机 (默认: 127.0.0.1)')
    parser.add_argument('--enable-forward', action='store_true', help='是否启用端口转发')

    return parser.parse_args()


def log(msg):
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}")


def relay(sock_in: socket.socket, sock_out: socket.socket):
    """单向转发数据，直到 EOF"""
    try:
        while data := sock_in.recv(BUFSIZE):
            sock_out.sendall(data)
    except OSError:
        pass
    finally:
        sock_in.close()
        sock_out.close()


def handle_client(client_conn: socket.socket, target_host: str, target_port: int):
    """接受一个客户端连接，并把它桥接到目标端口"""
    try:
        remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        remote.settimeout(60)
        remote.connect((target_host, target_port))
    except Exception as e:
        log(f"无法连接到目标 {target_host}:{target_port} -> {e}")
        client_conn.close()
        return

    # 启动两条线程，双向转发
    for args in [(client_conn, remote), (remote, client_conn)]:
        threading.Thread(target=relay, args=args, daemon=True).start()


def start_port_forwarding(listen_port, target_host, target_port):
    """启动端口转发服务"""
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        server.bind(("", listen_port))
        server.listen(100)
        log(f"启动端口转发: *:{listen_port}  ->  {target_host}:{target_port}")

        def forward_thread():
            try:
                while True:
                    conn, addr = server.accept()
                    log(f"新连接 {addr}")
                    threading.Thread(target=handle_client, args=(conn, target_host, target_port), daemon=True).start()
            except Exception as e:
                log(f"端口转发服务异常: {e}")
            finally:
                server.close()

        threading.Thread(target=forward_thread, daemon=True).start()
        return True
    except OSError as e:
        log(f"无法启动端口转发服务 {listen_port}: {e}")
        return False


def check_npx_availability():
    """检查npx命令是否可用（使用更可靠的方法，模仿CMD执行）"""
    cmd = ['cmd.exe', '/c', 'npx --version'] if os.name == 'nt' else ['npx', '--version']

    try:
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=10)
        if result.returncode == 0:
            version = result.stdout.strip()
            print(f"✅ 已找到 npx 命令 (版本: {version})" if os.name == 'nt' else "✅ 已找到 npx 命令")
            return True
    except Exception as e:
        print(f"❌ 检查 npx 失败: {e}" if os.name == 'nt' else f"❌ 检查 npx 失败: {e}")

    print(
        "⚠️  找不到 npx 命令。请确保已安装 Node.js 并添加到系统 PATH 中。" if os.name == 'nt' else "⚠️  找不到 npx 命令，请确保 Node.js 已安装。")
    print("👉 下载地址: https://nodejs.org")
    return False


def get_chromium_executable_path():
    """使用 Playwright 获取 Chromium 可执行路径"""
    try:
        with sync_playwright() as p:
            if not (exe_path := p.chromium.executable_path):
                print("❌ Playwright 未返回 Chromium 路径，请运行: npx playwright install chromium")
                return None
            if not (exe := Path(exe_path)).exists():
                print(f"❌ Chromium 可执行文件不存在: {exe_path}")
                return None
            print(f"✅ 成功获取 Chromium 路径: {exe_path}")
            return str(exe_path)
    except Exception as e:
        print(f"❌ 获取 Chromium 路径失败: {e}")
        return None


def get_screen_size():
    """获取屏幕尺寸"""
    try:
        import tkinter as tk
        root = tk.Tk()
        root.withdraw()
        size = (root.winfo_screenwidth(), root.winfo_screenheight())
        root.destroy()
        return size
    except Exception:
        return 1920, 1080


def kill_existing_browser_processes():
    """关闭已存在的浏览器进程，返回关闭的进程数量"""
    try:
        import psutil

        browser_names = ['chrome.exe', 'chromium.exe', 'chrome', 'chromium', 'google-chrome', 'google-chrome-stable']
        killed_count = 0

        print("🔍 正在查找并关闭已存在的浏览器进程...")

        for proc in psutil.process_iter(['pid', 'name']):
            try:
                if any(name in proc.info['name'].lower() for name in browser_names):
                    print(f"🛑 发现浏览器进程: {proc.info['name']} (PID: {proc.info['pid']})")
                    proc.kill()
                    proc.wait(timeout=5)
                    print(f"✅ 已关闭进程: {proc.info['name']} (PID: {proc.info['pid']})")
                    killed_count += 1
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
            except Exception as e:
                print(f"⚠️  处理进程时出错: {e}")

        print(f"✅ 共关闭了 {killed_count} 个浏览器进程" if killed_count else "✅ 未发现需要关闭的浏览器进程")
        return killed_count

    except ImportError:
        print("⚠️  未安装psutil库，无法自动关闭浏览器进程\n💡 请手动关闭所有浏览器窗口，或运行: pip install psutil")
        return 0
    except Exception as e:
        print(f"❌ 关闭浏览器进程时出错: {e}")
        return 0


def start_browser_and_mcp_instance(instance_id, args, chromium_path, screen_width, screen_height, mcp_port, relay_port,
                                   browser_debug_port):
    """启动单个浏览器和 MCP 实例"""
    try:
        # 创建用户数据目录
        user_data_dir = os.path.join(DATA_PATH,f"user_data_{instance_id}")
        os.makedirs(user_data_dir, exist_ok=True)

        # 构建浏览器启动参数
        browser_args = [
            chromium_path,
            f'--remote-debugging-port={browser_debug_port}',
            f'--user-data-dir={user_data_dir}',
            '--start-maximized',  # 启动时最大化窗口
            '--no-first-run', '--no-default-browser-check',
            '--disable-blink-features=AutomationControlled',
            '--disable-features=IsolateOrigins,site-per-process',
            '--disable-site-isolation-trials', '--disable-web-security',
            '--disable-features=BlockInsecurePrivateNetworkRequests',
            '--disable-features=InsecurePrivateNetworkRequestsSendPreflights',
            '--allow-running-insecure-content', '--ignore-certificate-errors',
        ]

        print(f"🚀 实例 {instance_id}: 启动浏览器...")
        browser_process = subprocess.Popen(browser_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        print(f"✅ 实例 {instance_id}: 浏览器启动成功 (PID: {browser_process.pid})")

        time.sleep(3)  # 等待浏览器启动

        # 启动 MCP 服务 - 使用与重构前相同的参数格式
        print(f"🚀 实例 {instance_id}: 启动 MCP 服务...")

        # 获取屏幕尺寸用于viewport参数
        width, height = get_screen_size()
        cdp_endpoint = f"http://127.0.0.1:{browser_debug_port}"

        mcp_args = [
            'npx', args.mcp_version,
            '--port', str(mcp_port),
            '--host', '127.0.0.1',
            '--user-data-dir', str(user_data_dir),
            '--isolated',
            '--viewport-size', f'{width},{height}',
            '--cdp-endpoint', cdp_endpoint,
            '--allowed-hosts', '*'
        ]

        if os.name == 'nt':
            # Windows系统下使用cmd启动新窗口
            cmd_line = subprocess.list2cmdline(mcp_args)
            print(f"🔧 启动 MCP: 端口={mcp_port} → CDP={cdp_endpoint}")
            mcp_process = subprocess.Popen([
                'cmd', '/c', 'start', '', 'cmd', '/k', cmd_line
            ])
            time.sleep(1)  # 小延迟避免窗口重叠
        else:
            print(f"🔧 启动 MCP: 端口={mcp_port} → CDP={cdp_endpoint}")
            mcp_process = subprocess.Popen(mcp_args)

        print(f"✅ 实例 {instance_id}: MCP 服务启动成功 (PID: {mcp_process.pid})")

        # 启动端口转发
        print(f"🚀 实例 {instance_id}: 启动端口转发...")
        relay_process = subprocess.Popen([
            sys.executable, __file__, '--relay-mode',
            '--local-port', str(mcp_port), '--remote-port', str(relay_port),
            '--instance-id', str(instance_id)
        ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        print(f"✅ 实例 {instance_id}: 端口转发启动成功 (PID: {relay_process.pid})")

        return browser_process, mcp_process, relay_process

    except Exception as e:
        print(f"❌ 实例 {instance_id}: 启动失败: {e}")
        return None, None, None


def start_mcp_instance(port, host, user_id, mcp_version, base_dir, cdp_endpoint):
    """启动单个 MCP 实例（前台保留窗口）"""

    # 使用配置文件中的DATA_PATH作为基础路径
    user_data_dir = os.path.join(DATA_PATH, 'user_data', f'user_{user_id}')
    os.makedirs(user_data_dir, exist_ok=True)

    print(f"📁 mcp使用用户数据目录: {user_data_dir}")

    width, height = get_screen_size()

    args = [
        'npx', mcp_version,
        '--port', str(port),
        '--host', host,
        '--user-data-dir', user_data_dir,
        '--isolated',
        '--viewport-size', f'{width},{height}',
        '--cdp-endpoint', cdp_endpoint,
        '--allowed-hosts', '*'
    ]

    if os.name == 'nt':
        cmd_line = subprocess.list2cmdline(args)
        print(f"🔧 启动 MCP: 端口={port} → CDP={cdp_endpoint}")
        process = subprocess.Popen([
            'cmd', '/c', 'start', '', 'cmd', '/k', cmd_line
        ])
        time.sleep(1)  # 小延迟避免窗口重叠
        return process
    else:
        print(f"🔧 启动 MCP: 端口={port} → CDP={cdp_endpoint}")
        return subprocess.Popen(args)


def main():
    """主函数"""
    print("🚀 开始批量启动 Chromium 和 MCP 实例")

    # 解析命令行参数
    args = parse_arguments()

    # 检查 npx 是否可用
    if not check_npx_availability():
        print("❌ npx 不可用，请确保已安装 Node.js 和 npm")
        return 1

    # 获取 Chromium 路径
    chromium_path = get_chromium_executable_path()
    if not chromium_path:
        print("❌ 无法获取 Chromium 路径")
        return 1

    # 获取屏幕尺寸
    screen_width, screen_height = get_screen_size()
    print(f"📺 屏幕尺寸: {screen_width}x{screen_height}")

    # 如果需要，关闭已存在的浏览器进程
    if args.kill_existing:
        kill_existing_browser_processes()

    # 启动所有实例
    all_processes = []

    for i in range(args.count):
        print(f"\n🎯 启动实例 {i + 1}/{args.count}")

        # 计算端口并启动实例
        browser_debug_port = args.base_browser_port + i
        mcp_port = args.base_mcp_port + i
        relay_port = args.base_relay_port + i

        browser_proc, mcp_proc, relay_proc = start_browser_and_mcp_instance(
            i, args, chromium_path, screen_width, screen_height,
            mcp_port, relay_port, browser_debug_port
        )

        if browser_proc and mcp_proc and relay_proc:
            all_processes.extend([browser_proc, mcp_proc, relay_proc])
            print(f"✅ 实例 {i + 1} 启动成功")
            print(f"   浏览器调试端口: {browser_debug_port}")
            print(f"   MCP 端口: {mcp_port}")
            print(f"   转发端口: {relay_port}")
        else:
            print(f"❌ 实例 {i + 1} 启动失败")

    # 等待所有进程
    if all_processes:
        print(f"\n⏳ 等待 {len(all_processes)} 个进程运行...")
        print("💡 按 Ctrl+C 可以停止所有进程")

        try:
            while all_processes:  # 等待任意一个进程结束
                for proc in all_processes[:]:
                    if proc.poll() is not None:
                        print(f"⚠️  进程 {proc.pid} 已结束")
                        all_processes.remove(proc)
                        break
                time.sleep(1)

        except KeyboardInterrupt:
            print("\n🛑 用户中断，正在停止所有进程...")
            for proc in all_processes:
                try:
                    proc.terminate()
                    proc.wait(timeout=5)
                except:
                    proc.kill()
            print("✅ 所有进程已停止")

    print("\n🏁 所有进程已结束")
    return 0


if __name__ == "__main__":
    main()
